#!/usr/bin/env python

import rospy
from nav_msgs.msg import Odometry
from geometry_msgs.msg import (
    Pose,
    Twist,
    Transform,
    TransformStamped,
    PoseWithCovarianceStamped,
)
from gazebo_msgs.msg import LinkStates
from std_msgs.msg import Header
import numpy as np
import math
import tf2_ros


class OdometryNode:
    # Set publishers
    pub_odom = rospy.Publisher("/odom", Odometry, queue_size=1)

    def __init__(self):
        # init internals
        self.last_received_pose = Pose()
        self.last_received_twist = Twist()
        self.last_recieved_stamp = None

        # Set the update rate
        rospy.Timer(rospy.Duration(0.05), self.timer_callback)  # 20hz

        self.tf_pub = tf2_ros.TransformBroadcaster()

        # Set subscribers
        rospy.Subscriber(
            "/amcl_pose", PoseWithCovarianceStamped, self.sub_robot_pose_update
        )
        rospy.Subscriber("/gazebo/link_states", LinkStates, self.sub_robot_pose2_update)

    def sub_robot_pose_update(self, msg):
        self.last_received_pose = msg.pose
        self.last_recieved_stamp = msg.header.stamp

    def sub_robot_pose2_update(self, msg):
        # Find the index of the racecar
        try:
            arrayIndex = msg.name.index(
                "racecar::base_footprint"
            )  # base_link -> base_footprint
        except ValueError as e:
            # Wait for Gazebo to startup
            pass
        else:
            # Extract our current position information
            self.last_received_twist = msg.twist[arrayIndex]

    def timer_callback(self, event):
        if self.last_recieved_stamp is None:
            return

        cmd = Odometry()
        cmd.header.stamp = self.last_recieved_stamp
        cmd.header.frame_id = "map"
        cmd.child_frame_id = (
            "base_footprint"  # This used to be odom, base_link -> base_footprint
        )
        cmd.pose = self.last_received_pose
        self.pub_odom.publish(cmd)

        tf = TransformStamped(
            header=Header(frame_id=cmd.header.frame_id, stamp=cmd.header.stamp),
            child_frame_id=cmd.child_frame_id,
            transform=Transform(
                translation=cmd.pose.pose.position, rotation=cmd.pose.pose.orientation
            ),
        )
        self.tf_pub.sendTransform(tf)


# Start the node
if __name__ == "__main__":
    rospy.init_node("amcl_odometry_node")
    node = OdometryNode()
    rospy.spin()
